# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/4 下午5:39
@Author yzpang
@Function: 
"""

import numpy as np
from modelserver.configs.server_config import MULTI_CLASS, MULTI_LABEL


def preprocess_function(examples, tokenizer, max_length, is_test=False, classification_type=MULTI_CLASS, labels_nums=2):
    """
    数据集预处理方法, 使用分词器将每一条样本分词处理
    :param labels_nums: 分类类别数量(层次分类时使用)
    :param classification_type: 分类类型(多分类/层次分类)
    :param examples:  样本
    :param tokenizer: 分词器
    :param max_length: 文本最大长度
    :param is_test: 是否测试使用, true不用返回标签
    :return:
    """
    result = tokenizer(examples['text'], max_length=max_length, truncation=True)
    if not is_test:
        if classification_type == MULTI_LABEL:
            result['labels'] = [float(1) if i in examples['label'] else float(0) for i in range(labels_nums)]
        else:
            result['labels'] = np.array([examples['label']], dtype=np.int64)
    return result


def read_local_dataset(path, label2id=None, is_test=False, classification_type=MULTI_CLASS):
    """加载数据集"""
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            if is_test:
                sentence = line.strip()
                yield {"text": sentence}
            else:
                items = line.strip().split('\t')
                if len(items) == 0:
                    continue
                if label2id is None:
                    if classification_type == MULTI_LABEL:
                        label = [l for l in items[-1].split(',')]
                    else:
                        label = items[-1]
                    yield {"text": "".join(items[:-1]), "label": label}
                else:
                    if classification_type == MULTI_LABEL:
                        label = [label2id[l] for l in items[-1].split(',')]
                    else:
                        label = label2id[items[-1]]
                    yield {"text": "".join(items[:-1]), "label": label}


def parse_labels(path):
    """解析标签列表"""
    id2label, label2id = {}, {}
    with open(path, 'r', encoding='utf-8') as f:
        for i, line in enumerate(f):
            l = line.strip()
            id2label[i] = l
            label2id[l] = i
    return id2label, label2id
